package com.alexander.scanregisterarouter.transform;

import com.alexander.scanregisterarouter.Util;
import com.android.build.api.transform.DirectoryInput;
import com.android.build.api.transform.Format;
import com.android.build.api.transform.JarInput;
import com.android.build.api.transform.QualifiedContent;
import com.android.build.api.transform.Transform;
import com.android.build.api.transform.TransformException;
import com.android.build.api.transform.TransformInput;
import com.android.build.api.transform.TransformInvocation;
import com.android.build.api.transform.TransformOutputProvider;
import com.android.build.gradle.internal.pipeline.TransformManager;

import org.apache.commons.io.FileUtils;
import org.apache.commons.io.IOUtils;
import org.gradle.internal.Pair;

import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.net.MalformedURLException;
import java.net.URL;
import java.net.URLClassLoader;
import java.util.Collection;
import java.util.Enumeration;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.jar.JarEntry;
import java.util.jar.JarFile;
import java.util.jar.JarOutputStream;
import java.util.zip.ZipEntry;

import javassist.CannotCompileException;
import javassist.ClassPool;
import javassist.CtClass;
import javassist.CtMethod;
import javassist.CtNewMethod;
import javassist.NotFoundException;

import static com.alexander.scanregisterarouter.Statics.AROUTER_INTERFACE_PATH;
import static com.alexander.scanregisterarouter.Statics.AROUTER_LOADPATH_METHOD_NAME;
import static com.alexander.scanregisterarouter.Statics.CLASS_FILE_SUFFIX;
import static com.alexander.scanregisterarouter.Statics.JAR_TMP_SUFFIX;
import static com.alexander.scanregisterarouter.Statics.PATH_FILE_NAME;
import static com.alexander.scanregisterarouter.Statics.ROUTER_MANAGER_NAME;
import static com.alexander.scanregisterarouter.Statics.ROUTER_MANAGER_REWRITE_METHOD_NAME;

/**
 * author : cjw
 * date : 2022/2/9
 * description :
 */
public class ScanRegisterARouterTransform extends Transform {
    private static final String TAG = "ScanRegisterARouterTransform";
    private Set<String> directoryInputClasses = new HashSet<>();
    private Map<String, String> routerTable = new HashMap<>();
    private Set<URL> needLoadClassOrJarList = new HashSet<>();
    // URLClassLoad只加载目录，loadClass或Class.forName()时，只接收'.'分割的全类名
    // /usr/app/                     com.alexander.test.MyApp
    private ClassLoader classLoader;
    // RouterManager类所处的jar包路径和全类名
    private Pair<JarInput, String> routerManagerJarClass;
    // 包含RouterManager类jar文件的新旧两个file对象
    private Pair<File, File> routerManagerParentJarData;
    // 非目标jar文件数据，即需要原封不动传递给下一个transform的jar文件
    private Map<File, File> notTargetJarData = new HashMap<>();
    private ClassPool classPool = ClassPool.getDefault();

    @Override
    public String getName() {
        return TAG;
    }

    @Override
    public Set<QualifiedContent.ContentType> getInputTypes() {
        return TransformManager.CONTENT_CLASS;
    }

    @Override
    public Set<? super QualifiedContent.Scope> getScopes() {
        return TransformManager.SCOPE_FULL_PROJECT;
    }

    @Override
    public boolean isIncremental() {
        return false;
    }

    @Override
    public void transform(TransformInvocation transformInvocation) throws TransformException, InterruptedException, IOException {
        super.transform(transformInvocation);
        System.out.println(TAG + "->start execute ScanRegisterARouterTransform");
        Collection<TransformInput> transformInputs = transformInvocation.getInputs();
        TransformOutputProvider outputProvider = transformInvocation.getOutputProvider();
        if (outputProvider != null) {
            outputProvider.deleteAll();
        }
//        System.out.println(TAG + "->transformInputs size:" + transformInputs.size());
        for (TransformInput input : transformInputs) {
            // 此transform被配置只在app模块中运行，因为app只是壳工程，所以app模块依赖其他所有业务模块，所以业务模块在
            // app模块内是以aar(jar)的方式存在的，所以，当transform运行的时候，DirectoryInput、JarInput两种都需要
            // 被扫描。
            // JarInput包含当前transform运行时所处的module所以来的所有第三方库的代码
            // DirectoryInput只包含当前transform运行时所处的module内的文件夹
            // 优先处理JarInput（因为有一些数据结构存在于jar中）
            Collection<JarInput> jarInputs = input.getJarInputs();
//            System.out.println(TAG + "->jarInputs size:" + jarInputs.size());
            // 首先找到ARouterPath接口class文件并加载、找到RouterManager类文件并加载；同时过滤掉非目标jar文件和
            // 非目标jar文件，并生成对应的contentLocation
            for (JarInput jarInput : jarInputs) {
                // jar文件如果不是目标对象（即是否需要修改其内容）则需要原封不动的传递给下一个transform，
                // 否则则需要把修改后的jar文件传递给下一个transform
                boolean targetJar = false;
                // 通过 把修改完的jar/不需要修改的jar文件拷贝至contentLocation 这种方式，
                // 把文件传递给下一个transform（必须有这个操作，否则打出来的包可能就不包含这个jar文件）
                File contentLocation = transformInvocation.getOutputProvider().getContentLocation(
                        jarInput.getName(), jarInput.getContentTypes(), jarInput.getScopes(),
                        Format.JAR);
//                System.out.println(TAG + "->jarInput name:" + jarInput.getName() + "，path：" +
//                        jarInput.getFile().getAbsolutePath());
                String jarPath = jarInput.getFile().getAbsolutePath();
                JarFile jarFile = new JarFile(jarPath);
                Enumeration<JarEntry> entries = jarFile.entries();
                // 遍历entries，找到对应的classes
                while (entries.hasMoreElements()) {
                    JarEntry jarEntry = entries.nextElement();
                    String name = jarEntry.getName();
                    boolean a = name.endsWith(CLASS_FILE_SUFFIX);
                    if (a && name.contains(ROUTER_MANAGER_NAME)) {
                        // 当前jarEntry是RouterManager,需要修改此jarEntry
                        System.out.println(TAG + "->RouterManager的类是：" + name + "，jar路径：" + jarPath);
                        routerManagerJarClass = Pair.of(jarInput, name);
                        // 记录目标jar文件的原始地址和contentLocation
                        targetJar = true;
                        routerManagerParentJarData = Pair.of(jarInput.getFile(), contentLocation);
                    } else if (a && name.contains(AROUTER_INTERFACE_PATH)) {
                        System.out.println(TAG + "->jar路径：" + jarPath + "，ARouterPath接口文件：" + name);
                        // 解压jar, 加载所有的class文件，为后面的反射运行loadPath方法做准备
                        final File workDir = File.createTempFile("unJar", "",
                                new File(System.getProperty("user.dir")));
                        workDir.delete();
                        workDir.mkdirs();
                        if (!workDir.isDirectory()) {
                            System.err.println(TAG + "->[0]mkdirs failed to create " + workDir);
                        } else {
                            System.out.println(TAG + "->[0]tempDir is：" + workDir.getAbsolutePath());
                        }
                        unJar(jarInput.getFile(), workDir);
                        // 解压完毕，设置classPath，加载jar中的所有class
                        needLoadClassOrJarList.add(new File(workDir.getAbsolutePath().concat(File.separator)).toURL());
                        System.out.println(TAG + "->success find ARouterPath interface");
                    }
                }
                // 记录非目标jar文件的原始地址和contentLocation
                if (!targetJar) {
                    notTargetJarData.put(jarInput.getFile(), contentLocation);
                }
            }
            // 然后扫描所有的 ARouter$Path$$ 类并加载, 反射实例化然后运行获取PathRouterTable
            for (JarInput jarInput : jarInputs) {
//                System.out.println(TAG + "->jarInput name:" + jarInput.getName() + "，path：" +
//                        jarInput.getFile().getAbsolutePath());
                String jarPath = jarInput.getFile().getAbsolutePath();
                try {
                    JarFile jarFile = new JarFile(jarPath);
                    Enumeration<JarEntry> entries = jarFile.entries();
                    // 遍历entries，找到对应的classes
                    while (entries.hasMoreElements()) {
                        JarEntry jarEntry = entries.nextElement();
                        String name = jarEntry.getName();
                        boolean a = name.endsWith(CLASS_FILE_SUFFIX);
                        if (a && name.contains(PATH_FILE_NAME)) {
                            System.out.println(TAG + "->jar路径：" + jarPath + "，class文件：" + name);
                            // 构造URLClassLoader需要的全类名
                            String classPkg = name.replaceAll(File.separator, "\\.")
                                    .replaceAll("\\.class", "");
                            // 解压jar, 加载class，反射运行loadPath方法
                            final File workDir = File.createTempFile("unJar", "",
                                    new File(System.getProperty("user.dir")));
                            workDir.delete();
                            workDir.mkdirs();
                            if (!workDir.isDirectory()) {
                                System.err.println(TAG + "->[1]mkdirs failed to create " + workDir);
                            } else {
                                System.out.println(TAG + "->[1]tempDir is：" + workDir.getAbsolutePath());
                            }
                            unJar(jarInput.getFile(), workDir);
                            // 解压完毕，设置classPath，加载jar中的所有class
                            needLoadClassOrJarList.add(new File(workDir.getAbsolutePath()
                                    .concat(File.separator)).toURL());
                            classLoader = new URLClassLoader(needLoadClassOrJarList
                                    .toArray(new URL[0]));
                            Thread.currentThread().setContextClassLoader(classLoader);
                            // 反射实例化并反射调用方法获取此类文件中的路由表单
                            Class<?> clazz = Class.forName(classPkg, true, classLoader);
                            Method loadPathMethod = clazz.getDeclaredMethod(AROUTER_LOADPATH_METHOD_NAME);
                            Object object = clazz.newInstance();
                            Map<String, String> tampTable = (Map<String, String>) loadPathMethod.invoke(object);
                            for (Map.Entry<String, String> entry : tampTable.entrySet()) {
                                routerTable.put(entry.getKey(), entry.getValue());
                            }
                        }
                    }
                } catch (MalformedURLException | ClassNotFoundException | NoSuchMethodException
                        | InstantiationException | IllegalAccessException | InvocationTargetException e) {
                    System.err.println(TAG + "->JarInput handle exception");
                    e.printStackTrace();
                }
            }
            // 处理DirectoryInput()
            Collection<DirectoryInput> directoryInputs = input.getDirectoryInputs();
//            System.out.println(TAG + "->directoryInputs size:" + directoryInputs.size());
            for (DirectoryInput dirInput : directoryInputs) {
                File dirFile = dirInput.getFile();
                if (dirFile != null && dirFile.exists()) {
//                    System.out.println(TAG + "->dirFile文件：" + dirFile.getAbsolutePath());
                    List<File> tmp = Util.getAllFile(dirFile.getAbsolutePath(), false);
                    tmp.forEach(file -> {
                        boolean a = file.getPath().endsWith(CLASS_FILE_SUFFIX);
                        boolean b = file.isFile();
                        boolean c = file.getName().startsWith(PATH_FILE_NAME);
                        if (a && b && c) {
                            // 构造URLClassLoader需要的全类名
                            String className = file.getAbsolutePath().replace(dirFile.getAbsolutePath()
                                    .concat(File.separator), "");
                            String classPkg = className.replaceAll(File.separator, "\\.")
                                    .replaceAll("\\.class", "");
                            directoryInputClasses.add(classPkg);
                        }
                    });
                    // 不为空说明此文件夹下有我们需要的 ARouter$$Path$$ 类文件
                    if (!directoryInputClasses.isEmpty()) {
                        needLoadClassOrJarList.add(dirFile.toURL());
                    }
                }
                // 通过 把class文件夹拷贝至contentLocation 这种方式，把类传递给下一个transform（必须有这个操作，
                // 否则打出来的包可能就不包含文件夹下的所有类）
                File contentLocation = transformInvocation.getOutputProvider().getContentLocation(
                        dirInput.getName(), dirInput.getContentTypes(), dirInput.getScopes(),
                        Format.DIRECTORY);
                FileUtils.copyDirectory(dirFile, contentLocation);
            }
            try {
                // 处理从DirectoryInput扫描到的所有 ARouter$$Path$$ 类
                // 首先构造classLoader加载所有需要的类
                classLoader = new URLClassLoader(needLoadClassOrJarList
                        .toArray(new URL[0]));
                Thread.currentThread().setContextClassLoader(classLoader);
                Iterator<String> iterator = directoryInputClasses.iterator();
                while (iterator.hasNext()) {
                    String element = iterator.next();
                    // 反射实例化并反射调用方法获取此类文件中的路由表单
                    Class<?> clazz = Class.forName(element, true, classLoader);
                    Method loadPathMethod = clazz.getDeclaredMethod(AROUTER_LOADPATH_METHOD_NAME);
                    Object object = clazz.newInstance();
                    Map<String, String> tampTable = (Map<String, String>) loadPathMethod.invoke(object);
                    for (Map.Entry<String, String> entry : tampTable.entrySet()) {
                        if (!routerTable.containsKey(entry.getKey())) {
                            routerTable.put(entry.getKey(), entry.getValue());
                        }
                    }
                }
            } catch (ClassNotFoundException | NoSuchMethodException | InstantiationException
                    | IllegalAccessException | InvocationTargetException e) {
                System.err.println(TAG + "->DirectoryInput handle exception");
                e.printStackTrace();
            }
        }
        System.out.println(TAG + "最终结果");
        System.out.println(TAG + "->RouterManager的类是：" + routerManagerJarClass.left
                + "，jar路径：" + routerManagerJarClass.right);
        for (Map.Entry<String, String> entry : routerTable.entrySet()) {
            System.err.println(TAG + "->路由表单：" + entry.getKey() + "<->" + entry.getValue());
        }
        try {
            // 根据路由表，给RouterManager.class添加一个新函数
            File newJarFile = fixClassFileInJar(routerManagerJarClass.left, routerManagerJarClass.right, routerTable);
            System.out.println(TAG + "->修改后的jar文件路径是：" + newJarFile.getAbsolutePath());
            // 把修改后的目标jar传递给下一个transform
            FileUtils.copyFile(newJarFile, routerManagerParentJarData.right);
        } catch (NotFoundException | CannotCompileException e) {
            System.err.println(TAG + "->fixClassFileInJar error");
            e.printStackTrace();
        }
        // 把未修改的非目标jar文件传递给下一个transform
        for (Map.Entry<File, File> entry : notTargetJarData.entrySet()){
            FileUtils.copyFile(entry.getKey(), entry.getValue());
        }
    }

    // 解压jar包
    private void unJar(File jarFile, File toDir) throws IOException {
        JarFile jar = new JarFile(jarFile);
        try {
            Enumeration entries = jar.entries();
            while (entries.hasMoreElements()) {
                JarEntry entry = (JarEntry) entries.nextElement();
                if (!entry.isDirectory()) {
                    InputStream in = jar.getInputStream(entry);
                    try {
                        File file = new File(toDir, entry.getName());
                        if (!file.getParentFile().mkdirs()) {
                            if (!file.getParentFile().isDirectory()) {
                                throw new IOException(
                                        "Mkdirs failed to create "
                                                + file.getParentFile()
                                                .toString());
                            }
                        }
                        OutputStream out = new FileOutputStream(file);
                        try {
                            byte[] buffer = new byte[8192];
                            int i;
                            while ((i = in.read(buffer)) != -1) {
                                out.write(buffer, 0, i);
                            }
                        } finally {
                            out.close();
                        }
                    } finally {
                        in.close();
                    }
                }
            }
        } finally {
            jar.close();
        }
    }

    /**
     * 修改jar包内的class类文件
     * @param jarInput
     * @param className
     * @param route
     * @return 修改后的jar的file对象
     * @throws NotFoundException
     * @throws IOException
     * @throws CannotCompileException
     */
    private File fixClassFileInJar(JarInput jarInput, String className, Map<String, String> route)
            throws NotFoundException, IOException, CannotCompileException {
        String classPkg = className.replaceAll(File.separator, "\\.")
                .replaceAll("\\.class", "");
        classPool.appendClassPath(jarInput.getFile().getAbsolutePath());
        CtClass ctClass = classPool.getCtClass(classPkg);
        if (ctClass.isFrozen()) {
            ctClass.defrost();
        }
        // 首先获取到方法对象
        CtMethod ctMethod = ctClass.getMethod(ROUTER_MANAGER_REWRITE_METHOD_NAME, "(Ljava/util/Map;)V");
        System.out.println(TAG + "->方法为：" + ctMethod.getName());
        // 构造需要插入的代码
        char doubleQuotationMark = '"';
        StringBuilder newMethodBody = new StringBuilder();
        for (Map.Entry<String, String> entry : route.entrySet()) {
            newMethodBody.append("router.put(")
                    .append(doubleQuotationMark)
                    .append(entry.getKey())
                    .append(doubleQuotationMark)
                    .append(", ")
                    .append(doubleQuotationMark)
                    .append(entry.getValue())
                    .append(doubleQuotationMark)
                    .append(");")
                    .append("\n");
        }
        System.out.println(TAG + "->代码为：" + newMethodBody.toString());
        // 把代码插入方法的前边(即从方法体第一行开始插入)
        ctMethod.insertBefore(newMethodBody.toString());
        // 获取被修改后的类的字节码
        byte[] fixedClassBytes = ctClass.toBytecode();
        ctClass.detach();
        // 当前jar的JarFile示例
        JarFile jarFile = new JarFile(jarInput.getFile());
        // 修改后的jar的路径和file实例（重新写入jar文件）
        String tmpJarPath = jarInput.getFile().getParent().concat(File.separator)
                .concat(jarInput.getName().concat(JAR_TMP_SUFFIX));
        File tmpJarFile = new File(tmpJarPath);
        JarOutputStream jos = new JarOutputStream(new FileOutputStream(tmpJarFile));
        Enumeration<JarEntry> entries = new JarFile(jarInput.getFile()).entries();
        while (entries.hasMoreElements()) {
            JarEntry jarEntry = entries.nextElement();
            jos.putNextEntry(new ZipEntry(jarEntry.getName()));
            // 找到被修改的那个jarEntry,把修改后的字节码写入
            if (jarEntry.getName().equals(className)) {
                System.out.println(TAG + "->向tmpJarFile中写入修改后的字节码");
                jos.write(fixedClassBytes);
            } else {
                InputStream is = jarFile.getInputStream(jarEntry);
                byte[] idleClassBytes = IOUtils.toByteArray(is);
                jos.write(idleClassBytes);
            }
            jos.closeEntry();
        }
        jos.close();
        jarFile.close();
        return tmpJarFile;
    }
}
